from output import Output
from plugin_system import Plugin
from parameter_system import Enum

class EnergyOutput(Output):
    def __init__(self, filename=None, freq=None):
        """
        Takes the filename of the file that the data
        will be written to, the frequency at which the
        data will be written
        """
        self.filename = filename
        self.fl = None
        self.frequency = freq

    def initialize(self):
        """
        Opens the output file.
        """
        self.fl = open(self.filename, "w")

    def finalize(self):
        """
        Closes the output file.
        """
        if self.fl == None:
            raise Exception, "Cannot finalize output that has not be initialized."
        self.fl.close()

    def output(self, system, step):
        if step % self.frequency == 0:
            # we expect a list of 3-tuples
            ke = system.kinetic_energy
            pe = system.potential_energy
            s = "Step: %s, KE: %s, PE: %s\n" % (step, ke, pe)
            self.fl.write(s)

class EnergyOutputPluginAdapter(Plugin, EnergyOutput):
    plugin_name = "EnergyOutput"
    parameter_definitions = {"filename" : (str, True, "File to which to write data."),
                             "freq" : (int, True, "Frequency (in steps) at which to write data.")}
    def __init__(self, *args, **user_params):
        params = self.configure_parameters(**user_params)
        EnergyOutput.__init__(self, filename=params["filename"], freq=params["freq"])
